import json
import os
import shutil
import random
 
labels = ['coal', 'bicycle', 'car', 'motorbike', 'aeroplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'sofa', 'pottedplant', 'bed', 'diningtable', 'toilet', 'tvmonitor', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
 
 
def generate_diretorys(wd):
 
 
    data_base_dir = os.path.join(wd, "VOCdevkit/")
    try:
        shutil.rmtree(data_base_dir)
        print("{} exist, will delete it".format(data_base_dir))
    except:
        print("{} does exist, does not need to delete!".format(data_base_dir))
    os.mkdir(data_base_dir)
 
    work_sapce_dir = os.path.join(data_base_dir, "VOC2007/")
    os.mkdir(work_sapce_dir)
 
    annotation_dir = os.path.join(work_sapce_dir, "Annotations/")
    os.mkdir(annotation_dir)
 
    image_dir = os.path.join(work_sapce_dir, "JPEGImages/")
    os.mkdir(image_dir)
 
    yolo_labels_dir = os.path.join(work_sapce_dir, "YOLOLabels/")
    os.mkdir(yolo_labels_dir)
 
    yolov5_images_dir = os.path.join(data_base_dir, "images/")
    os.mkdir(yolov5_images_dir)
 
    yolov5_labels_dir = os.path.join(data_base_dir, "labels/")
    os.mkdir(yolov5_labels_dir)
 
    yolov5_images_train_dir = os.path.join(yolov5_images_dir, "train/")
    os.mkdir(yolov5_images_train_dir)
 
    yolov5_images_test_dir = os.path.join(yolov5_images_dir, "val/")
    os.mkdir(yolov5_images_test_dir)
 
    yolov5_labels_train_dir = os.path.join(yolov5_labels_dir, "train/")
    os.mkdir(yolov5_labels_train_dir)
 
    yolov5_labels_test_dir = os.path.join(yolov5_labels_dir, "val/")
    os.mkdir(yolov5_labels_test_dir)
 
    print("创建VOC数据集格式多级目录完成\n")
 
 
 
 
def convert(jsonfile, yoloPath):
    """
    Args:
        jsonfile: anylabeling标注的结果文件 ，需要转化成coco128-seg数据集标注文件格式
        yoloPath:
    Returns:
    """
 
    assert os.path.exists(yoloPath), "'{}' does not exits".format(yoloPath)
    assert os.path.isfile(jsonfile), "'{}' does not exits".format(jsonfile)
 
    # 读取json标注文件
    with open(jsonfile, 'r') as load_file:
        load_dict = json.load(load_file)
 
    # 准备yolo标注输出文件
    fileName = os.path.split(jsonfile)[-1]
    file_id, extension = os.path.splitext(fileName)
    output_fileName = file_id + ".txt"
    outputPath = os.path.join(yoloPath, output_fileName)
    outputFile = open(outputPath, 'w')
    # print("\r jason文件名字：%s" %fileName)
 
    #json2yolo
    height = load_dict["imageHeight"]
    width = load_dict["imageWidth"]
    for item in load_dict["shapes"]:
        # print(item)
        label_int = labels.index(item["label"])
        row_str = ""
        row_str += str(label_int)
        if not item["points"]:
            continue
        for point in item["points"]:
            x = round(float(point[0])/width, 6) #小数保留6位
            y = round(float(point[1])/height, 6) #小数保留6位
            row_str += " " + str(x) + " " + str(y)
 
        row_str+="\n"
        outputFile.write(row_str)
    outputFile.close()
 
 

def add_background(inputdir,output_path,val_ratio):
    print("背景")
    img_list = os.listdir(inputdir)
    num = len(img_list)
    random.seed(0)
    eval_img_files = random.sample(img_list, k=int(num*val_ratio))

    for img_file in img_list:
        if img_file in eval_img_files:
            img_path = os.path.join(inputdir,img_file)
            val_img_path = os.path.join(output_path,"VOCdevkit","images","val")
            shutil.copy(img_path, val_img_path)
        else:
            img_path = os.path.join(inputdir,img_file)
            train_img_path = os.path.join(output_path,"VOCdevkit","images","train")
            shutil.copy(img_path, train_img_path)


 
 
def json2yolo(inputPath, output_path, val_ratio):
    """
    Args:
        inputPath: json标注文件与图片所在的目录
        output_path: 输出目录，会在该目录下生成VOC数据集格式多级目录
        val_ratio: 验证集占比
    Returns:
    """
    #
    assert os.path.exists(inputPath), "'{}' does not exist".format(inputPath)
    assert os.path.exists(output_path), "'{}' does not exist".format(output_path)
 
 
    generate_diretorys(wd=output_path)
 
    wd = output_path
    train_file = open(os.path.join(wd, "yolov5_train.txt"), 'w')
    val_file = open(os.path.join(wd, "yolov5_val.txt"), 'w')
    train_file.close()
    val_file.close()
    train_file = open(os.path.join(wd, "yolov5_train.txt"), 'a')
    val_file = open(os.path.join(wd, "yolov5_val.txt"), 'a')
 
 
 
 
    # 检查图片与标签文件一一对应
    raw_files = os.listdir(inputPath)
    json_files = [] # jason文件完整路径
    img_files = [] # 图片文件完整路径
    for file in raw_files:
        # print("file:",file)
        fileName, extension = os.path.splitext(file)
        # print("extension:",extension)
        # print(type(extension))
        if extension in ['.jpg','.jpeg','.ico','.tiff', '.png','.bmp']:
            json_file = fileName + ".json"
            json_path = os.path.join(inputPath, json_file)
            # print("json_path:",json_file)
 
            if os.path.exists(json_path):
                img_path = os.path.join(inputPath,file)
 
                json_files.append(json_path)
                img_files.append(img_path)
 
 
                another_img_path = os.path.join(wd,"VOCdevkit","VOC2007","JPEGImages")
                another_jason_path = os.path.join(wd,"VOCdevkit","VOC2007","Annotations")
                shutil.copy(json_path, another_jason_path)
                shutil.copy(img_path, another_img_path)
 
    num = len(json_files)
    random.seed(0)
    eval_json_files = random.sample(json_files, k=int(num*val_ratio))
 
    i = 0
    for json_path in json_files:
        # print("json_path:",json_path)
 
        another_yolo_path = os.path.join(wd,"VOCdevkit","VOC2007","YOLOLabels")
 
        json_file_names_info = os.path.split(json_path)[-1]
        json_file_name_id, ext= os.path.splitext(json_file_names_info)
        yolo_file_name = json_file_name_id + ".txt"
 
        index = json_files.index(json_path)
        img_file_path = img_files[index]
        img_file_name = img_file_path.split("/")[-1]
 
 
        # print("\r json_file_name_id:{}   yolo_file_name :{}".format(json_file_name_id,yolo_file_name))
 
        if json_path in eval_json_files:
            index = json_files.index(json_path)
            img_path = img_files[index]
 
 
            val_img_path = os.path.join(wd,"VOCdevkit","images","val")
            val_json_path = os.path.join(wd,"VOCdevkit","labels","val")
            yolo_file = os.path.join(val_json_path, yolo_file_name)
            
 
            convert(jsonfile=json_path, yoloPath=val_json_path)
            shutil.copy(img_path, val_img_path)
            shutil.copy(yolo_file, another_yolo_path)
            # shutil.copy(yolo_file, val_json_path)
 
 
            val_file.write(os.path.join(val_img_path, img_file_name) + '\n')
 
 
        else:
            index = json_files.index(json_path)
            img_path = img_files[index]
            train_img_path = os.path.join(wd,"VOCdevkit","images","train")
            train_json_path = os.path.join(wd,"VOCdevkit","labels","train")
            yolo_file = os.path.join(train_json_path, yolo_file_name)
            print("json_path：",json_path)
            
 
 
            convert(jsonfile=json_path, yoloPath=train_json_path)

            shutil.copy(img_path, train_img_path)
            shutil.copy(yolo_file, another_yolo_path)
            # print("train_json_path",train_json_path)
            # print("yolo_file",yolo_file)
            # shutil.copy(yolo_file, train_json_path)
 
 
 
            train_file.write(os.path.join(train_img_path, img_file_name) + '\n')
 
        i+=1
        print("\r processing [{}/{}]".format(i, num), end="")
 
    
if __name__ == "__main__":
    outputdir = r"F:/SY_data_set/皮带轻载，中载，重载/yolo-seg"
    if not os.path.exists(outputdir):
        os.makedirs(outputdir)
    inputdir = r"F:/SY_data_set/皮带轻载，中载，重载/yolo-seg1/VOCdevkit/VOC2007/JPEGImages"
    input_background_dir = r"F:\SY_data_set\皮带轻载，中载，重载\背景"
    json2yolo(inputPath=inputdir, output_path=outputdir, val_ratio=0.2)
    add_background(input_background_dir, outputdir, val_ratio=0.2)